import akshare as ak
import pandas as pd
import numpy as np
import talib


# 获取股票日线数据，增加重试机制
def get_stock_daily_data(symbol, start_date, end_date, retries=3):
    for _ in range(retries):
        try:
            stock_data = ak.stock_zh_a_daily(symbol=symbol, start_date=start_date, end_date=end_date)
            return stock_data
        except Exception as e:
            print(f"Error fetching data (attempt {_ + 1}): {e}")
    print("Failed to fetch data after multiple attempts.")
    return pd.DataFrame()


# 判断红三兵形态，优化条件判断
def is_red_three_soldiers(data):
    return data[(data['close'].shift(2) < data['open'].shift(2)) &
                (data['close'].shift(1) > data['open'].shift(1)) &
                (data['close'] > data['open']) &
                (data['close'].shift(1) < data['close']) &
                (data['close'].shift(2) < data['close'].shift(1))].index


# 获取缩倍量阴线数据，改进成交量变化倍数计算
def get_shrinking_doubled_yin_line_data(symbol, start_date, end_date):
    stock_data = get_stock_daily_data(symbol, start_date, end_date)
    if stock_data.empty:
        return []
    # 计算成交量的变化倍数（使用与前一日成交量的比值）
    stock_data['volume_change_ratio'] = stock_data['volume'].div(stock_data['volume'].shift(1)) - 1
    shrinking_doubled_yin_lines = []
    for i in range(1, len(stock_data)):
        if (stock_data['close'][i] < stock_data['open'][i]) and \
                (stock_data['volume_change_ratio'][i] <= -0.5) and \
                (stock_data['low'][i] > stock_data['low'][i - 1]):
            shrinking_doubled_yin_lines.append(stock_data['date'][i])
    return shrinking_doubled_yin_lines


# 获取MACD红柱二次启动数据
def get_macd_second_red_column_data(symbol, start_date, end_date):
    stock_data = get_stock_daily_data(symbol, start_date, end_date)
    if stock_data.empty:
        return []
    # 计算MACD指标
    close_prices = stock_data['close'].values
    macd, signal, hist = talib.MACD(close_prices)
    stock_data['macd'] = macd
    stock_data['signal'] = signal
    stock_data['hist'] = hist
    second_red_column_dates = []
    for i in range(1, len(stock_data)):
        if (hist[i] > 0) and (hist[i - 1] < 0):
            second_red_column_dates.append(stock_data['date'][i])
    return second_red_column_dates


# 主力筹码分布模型（简化示例，仅计算一定周期内收盘价的均值作为参考）
def analyze_chip_distribution(data):
    period = 20
    close_mean = data['close'].rolling(window=period).mean()
    return "筹码可能在集中" if data['close'].iloc[-1] > close_mean.iloc[-1] else "筹码可能在分散"


# 主力资金流向模型（结合价格变化和成交量变化判断）
def analyze_capital_flow(data):
    volume_change = data['volume'].pct_change()
    price_change = data['close'].pct_change()
    if volume_change.iloc[-1] > 0 and price_change.iloc[-1] > 0:
        return "可能有资金流入且价格上涨"
    elif volume_change.iloc[-1] > 0 and price_change.iloc[-1] < 0:
        return "可能有资金流入但价格下跌"
    elif volume_change.iloc[-1] < 0 and price_change.iloc[-1] > 0:
        return "可能有资金流出但价格上涨"
    else:
        return "可能有资金流出且价格下跌"


# 量价关系模型，增加更多情况判断
def analyze_price_volume_relation(data):
    price_change = data['close'].pct_change()
    volume_change = data['volume'].pct_change()
    if price_change.iloc[-1] > 0 and volume_change.iloc[-1] > 0:
        return "价升量增"
    elif price_change.iloc[-1] < 0 and volume_change.iloc[-1] < 0:
        return "价跌量缩"
    elif price_change.iloc[-1] > 0 and volume_change.iloc[-1] < 0:
        return "价升量缩"
    elif price_change.iloc[-1] < 0 and volume_change.iloc[-1] > 0:
        return "价跌量增"
    else:
        return "价格和成交量无明显变化"


# 成交量形态模型 - 高量柱判断
def detect_high_volume_column(data):
    max_volume = data['volume'].max()
    high_volume_dates = data[data['volume'] == max_volume]['date'].tolist()
    return high_volume_dates


# 三日爆量模型，优化参数和逻辑判断
def three_day_huge_volume(data):
    yesterday_volume = data['volume'].iloc[-2]
    today_volume = data['volume'].iloc[-1]
    if today_volume >= 3 * yesterday_volume:  # 可调整倍数参数
        next_two_days = data.iloc[-3:-1]
        if next_two_days['close'].min() > data['close'].iloc[-1] and \
                (next_two_days['close'].max() > data['high'].iloc[-1]):
            return "三日爆量模型成立"
    return "三日爆量模型不成立"


# 量能平台突破模型，优化判断逻辑
def volume_platform_breakthrough(data):
    max_volume_date = data['volume'].idxmax()
    max_volume_close = data.at[max_volume_date, 'close']
    after_max_volume_data = data[data.index > max_volume_date]
    for _, row in after_max_volume_data.iterrows():
        if row['close'] > max_volume_close and row['volume'] > data['volume'].mean():
            return "量能平台突破"
    return "量能平台未突破"


def main():
    start_date = "2023-01-01"
    end_date = "2024-10-24"
    all_stocks_info = ak.stock_info_a_code_name()
    qualified_stocks = []
    for _, row in all_stocks_info.iterrows():
        symbol = row['code']
        if symbol.startswith(("00", "300", "60")):
            prefix = "sz" if symbol.startswith(("00", "300")) else "sh"
            symbol_with_prefix = prefix + symbol
            stock_data = get_stock_daily_data(symbol_with_prefix, start_date, end_date)
            if not stock_data.empty:
                red_three_soldiers_dates = is_red_three_soldiers(stock_data)
                shrinking_doubled_yin_lines = get_shrinking_doubled_yin_line_data(symbol_with_prefix, start_date, end_date)
                macd_second_red_column_dates = get_macd_second_red_column_data(symbol_with_prefix, start_date, end_date)
                chip_distribution_result = analyze_chip_distribution(stock_data)
                capital_flow_result = analyze_capital_flow(stock_data)
                price_volume_relation_result = analyze_price_volume_relation(stock_data)
                high_volume_dates = detect_high_volume_column(stock_data)
                three_day_huge_volume_result = three_day_huge_volume(stock_data)
                volume_platform_breakthrough_result = volume_platform_breakthrough(stock_data)

                conditions = {
                    "红三兵形态": bool(len(red_three_soldiers_dates)),
                    "缩倍量阴线": bool(shrinking_doubled_yin_lines),
                    "MACD红柱二次启动": bool(macd_second_red_column_dates),
                    "筹码可能在集中": chip_distribution_result == "筹码可能在集中",
                    "可能有资金流入且价格上涨": "可能有资金流入且价格上涨" in capital_flow_result,
                    "价升量增": price_volume_relation_result == "价升量增",
                    "高量柱": bool(high_volume_dates),
                    "三日爆量模型成立": three_day_huge_volume_result == "三日爆量模型成立",
                    "量能平台突破": volume_platform_breakthrough_result == "量能平台突破"
                }
                if all(conditions.values()):
                    qualified_stocks.append(symbol)

    print("满足所有条件的股票列表:", qualified_stocks)


if __name__ == "__main__":
    main()